import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense,SimpleRNN,Dropout
import os
import pandas as pd
import math
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error,mean_absolute_error

maotai=pd.read_csv('./SH600519.csv')
training_set=maotai.iloc[0:2426-300,2:3].values #前[2426-300=2126]天的开盘价为训练集，表格从0开始计数，[2:3]前闭后开，提取开盘价那一列
test_set=maotai.iloc[2426-300,2:3].values #后300天的开盘价为训练集，表格从0开始计数，[2:3]前闭后开，提取开盘价那一列

#归一化
sc=MinMaxScaler(feature_range=(0,1)) #定义归一化到[0,1]之间
training_set_scaled=sc.fit_transform(training_set)
test_set_scaled=sc.fit_transform(test_set)

x_train=[]
y_train=[]

x_test=[]
y_test=[]

#以连续60天的开盘价格作为输入特征，第61天的数据作为标签,for循环并构建2426-300-60=2066组数据
for i in range(60,len(training_set_scaled)):
    x_train.append(training_set_scaled[i-60:i,0])
    y_train.append(training_set_scaled[i,0])

np.random.seed(7)
np.random.shuffle(x_train)
np.random.seed(7)
np.random.shuffle(y_train)
tf.random.set_seed(7)
#将list转为array,符合RNN输入要求：[送入样本数，循环时间展开步数，每个时间步输入特征个数]
x_train,y_train=np.array(x_train),np.array(y_train)
x_train=np.reshape(x_train, (x_train.shape[0],60,1))

#以连续60天的开盘价格作为输入特征，第61天的数据作为标签,for循环并构建300-60=240组数据
for i in range(60,len(test_set_scaled)):
    x_test.append(test_set_scaled[i-60:i,0])
    x_test.append(test_set_scaled[i,0])
#将list转为array,符合RNN输入要求：[送入样本数，循环时间展开步数，每个时间步输入特征个数]
x_test,y_test=np.array(x_test),np.array(y_test)
x_test=np.reshape(x_test, (x_test.shape[0],60,1))


model=tf.keras.Sequential([
    SimpleRNN(80,return_sequences=True),
    Dropout(0.2),
    SimpleRNN(100),
    Dropout(0.2),
    Dense(1)
])
model.compile(optimizer=tf.keras.optimizers.Adam(0.01),
              loss=tf.keras.losses.mean_squared_error)

# print(os.path.dirname(os.path.abspath("checkpoint/rnn_onhot.ckpt")))
checkpoint_save_path=os.path.abspath("checkpoint/rnn_stock.ckpt")
print(checkpoint_save_path)
if os.path.exists(checkpoint_save_path+'.index'):
    print('-------------------load the model-------------------')
    model.load_weights(checkpoint_save_path)

cp_callback=tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,
                                                save_weights_only=True,
                                                save_best_only=True,
                                                monitor='val_loss')       

history=model.fit(x_train,y_train,batch_size=64,epochs=50,validation_freq=1, validation_data=(x_test,y_test), callbacks=[cp_callback])
model.summary()

file=open('./weights_rnn_stock.txt','w')
for v in model.trainable_variables:
    file.write(str(v.name)+'\n')
    file.write(str(v.shape)+'\n')
    file.write(str(v.numpy())+'\n')
file.close()


#预测
# 对测试集合进行预测
predicted_stock_price=model.predict(x_test)
# 对预测数据还原，反归一化到原始范围
predicted_stock_price=sc.inverse_transform(predicted_stock_price)
# 对真实数据还原，反归一化到原始范围
real_stock_price=sc.inverse_transform(test_set[60:])

# preNum=int(input("输入测试字母数量："))
# for i in range(preNum):
#     alphabet1=input("输入测试字母：")
#     alphabet=[w_to_id[a] for a in alphabet1]
#     alphabet=np.reshape(alphabet, [1,4])
#     result=model.predict([alphabet])
#     pred=tf.argmax(result,axis=1)
#     pred=int(pred)
#     tf.print(alphabet1+'->'+input_word[pred])

#模型效果评价
mse=mean_squared_error(predicted_stock_price,real_stock_price)
rmse=math.sqrt(mse)
mae=mean_absolute_error(predicted_stock_price,real_stock_price)

print('均方误差：%.6f'%mse)

print('均方根误差：%.6f'%rmse)

print('平均绝对误差：%.6f'%mae)